{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "26134228",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "#全局变量\n",
    "hub_token = open('/root/hub_token.txt').read().strip()\n",
    "repo_id = 'lansinuote/diffusion.1.unconditional'\n",
    "push_to_hub = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8560ef3b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using custom data configuration huggan--flowers-102-categories-2ab3d0588f2a8da7\n",
      "Found cached dataset parquet (/root/.cache/huggingface/datasets/huggan___parquet/huggan--flowers-102-categories-2ab3d0588f2a8da7/0.0.0/2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e5500f4edbd145b2b63b63f4506d3a1a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Pushing split train to the Hub.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ed4ddd42675a4f6b8f176aca4eeecd8c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/9 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f0c60734f48c4703ae7b8e28e02f72b9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Pushing dataset shards to the dataset hub:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using custom data configuration lansinuote--diffusion.1.unconditional-a9e8dd8d646db18f\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading and preparing dataset None/None to /root/.cache/huggingface/datasets/lansinuote___parquet/lansinuote--diffusion.1.unconditional-a9e8dd8d646db18f/0.0.0/2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "84ccb0e10fed416cbafd9137bb0dc338",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data files:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d7f67a36a1e840a9b0ed538eb81a3659",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/347M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4691ddd0ee66405f955cbd37347d6845",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Extracting data files:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split:   0%|          | 0/8189 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset parquet downloaded and prepared to /root/.cache/huggingface/datasets/lansinuote___parquet/lansinuote--diffusion.1.unconditional-a9e8dd8d646db18f/0.0.0/2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec. Subsequent calls will reuse this data.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(Dataset({\n",
       "     features: ['image'],\n",
       "     num_rows: 8189\n",
       " }),\n",
       " {'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=752x500>})"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "\n",
    "def get_dataset():\n",
    "    #加载数据集\n",
    "    dataset = load_dataset('huggan/flowers-102-categories')\n",
    "    #dataset = load_dataset('huggan/pokemon')\n",
    "\n",
    "    return dataset\n",
    "\n",
    "\n",
    "if push_to_hub:\n",
    "    dataset = get_dataset()\n",
    "    dataset.push_to_hub(repo_id=repo_id, token=hub_token)\n",
    "\n",
    "#直接使用我处理好的数据集\n",
    "dataset = load_dataset(path=repo_id, split='train')\n",
    "\n",
    "dataset, dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "102e0a89",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Dataset({\n",
       "     features: ['image'],\n",
       "     num_rows: 8189\n",
       " }),\n",
       " {'image': tensor([[[-0.8118, -0.7882, -0.8039,  ..., -0.9686, -0.9216, -0.8353],\n",
       "           [-0.7647, -0.7647, -0.7804,  ..., -0.9451, -0.9451, -0.8667],\n",
       "           [-0.7333, -0.7412, -0.7647,  ..., -0.8902, -0.9294, -0.8824],\n",
       "           ...,\n",
       "           [-0.0510, -0.0510, -0.2314,  ..., -0.2627,  0.0196, -0.1765],\n",
       "           [-0.1686, -0.1294, -0.3176,  ..., -0.3490, -0.0196, -0.0196],\n",
       "           [-0.3020, -0.2706, -0.2078,  ..., -0.3882, -0.1451,  0.1765]],\n",
       "  \n",
       "          [[-0.7412, -0.7176, -0.7412,  ..., -0.9686, -0.9059, -0.7725],\n",
       "           [-0.6941, -0.6706, -0.6863,  ..., -0.9373, -0.9294, -0.8118],\n",
       "           [-0.6549, -0.6627, -0.6784,  ..., -0.8275, -0.9059, -0.8510],\n",
       "           ...,\n",
       "           [ 0.1608,  0.1529, -0.0118,  ..., -0.0431, -0.0902, -0.7020],\n",
       "           [ 0.0431,  0.0588, -0.0980,  ..., -0.1373, -0.0510, -0.3804],\n",
       "           [-0.0902, -0.0667, -0.0431,  ..., -0.2000, -0.1059, -0.0039]],\n",
       "  \n",
       "          [[-0.8980, -0.8824, -0.9059,  ..., -0.9843, -0.9765, -0.9529],\n",
       "           [-0.8510, -0.8431, -0.8667,  ..., -0.9608, -0.9765, -0.9529],\n",
       "           [-0.8275, -0.8431, -0.8667,  ..., -0.8980, -0.9529, -0.9451],\n",
       "           ...,\n",
       "           [-0.5922, -0.6863, -0.8196,  ..., -0.2863,  0.1843,  0.0980],\n",
       "           [-0.7412, -0.7333, -0.8118,  ..., -0.3882,  0.1137,  0.2863],\n",
       "           [-0.9137, -0.8902, -0.7569,  ..., -0.4510, -0.0824,  0.4510]]])})"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torchvision\n",
    "\n",
    "#图像增强\n",
    "compose = torchvision.transforms.Compose([\n",
    "    torchvision.transforms.Resize(\n",
    "        64, interpolation=torchvision.transforms.InterpolationMode.BILINEAR),\n",
    "    torchvision.transforms.RandomCrop(64),\n",
    "    #torchvision.transforms.RandomHorizontalFlip(),\n",
    "    torchvision.transforms.ToTensor(),\n",
    "    torchvision.transforms.Normalize([0.5], [0.5])\n",
    "])\n",
    "\n",
    "\n",
    "def f(data):\n",
    "    image = [compose(i) for i in data['image']]\n",
    "    return {'image': image}\n",
    "\n",
    "\n",
    "#因为图像增强在每个epoch中要动态计算,所以不能简单地用map处理\n",
    "dataset.set_transform(f)\n",
    "\n",
    "dataset, dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3ef5b1fc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(512,\n",
       " {'image': tensor([[[[ 0.3569, -0.2157, -0.7882,  ..., -0.6078, -0.7098, -0.6941],\n",
       "            [ 0.2078, -0.4745, -0.9137,  ..., -0.6157, -0.6392, -0.6706],\n",
       "            [ 0.0196, -0.6941, -0.9294,  ..., -0.6314, -0.6000, -0.5451],\n",
       "            ...,\n",
       "            [-0.3569, -0.5608, -0.9451,  ..., -0.9686, -0.9529, -0.8824],\n",
       "            [-0.5137, -0.8980, -0.9843,  ..., -0.9686, -0.9294, -0.8275],\n",
       "            [-0.8275, -0.9686, -0.9765,  ..., -0.9451, -0.8510, -0.7490]],\n",
       "  \n",
       "           [[ 0.2314, -0.2941, -0.7882,  ..., -0.5059, -0.5922, -0.5922],\n",
       "            [ 0.0667, -0.5294, -0.9059,  ..., -0.5137, -0.5137, -0.5608],\n",
       "            [-0.1137, -0.7176, -0.9294,  ..., -0.5137, -0.4353, -0.3725],\n",
       "            ...,\n",
       "            [-0.1608, -0.4824, -0.9451,  ..., -0.9686, -0.9529, -0.8745],\n",
       "            [-0.3882, -0.8745, -0.9843,  ..., -0.9686, -0.9137, -0.7961],\n",
       "            [-0.7804, -0.9686, -0.9765,  ..., -0.9451, -0.8118, -0.6941]],\n",
       "  \n",
       "           [[-0.1137, -0.5529, -0.8039,  ..., -0.7020, -0.7255, -0.7490],\n",
       "            [-0.3020, -0.7020, -0.9059,  ..., -0.6784, -0.6941, -0.7255],\n",
       "            [-0.4510, -0.8039, -0.9216,  ..., -0.7020, -0.6392, -0.5843],\n",
       "            ...,\n",
       "            [-0.7961, -0.7961, -0.9686,  ..., -0.9686, -0.9529, -0.8980],\n",
       "            [-0.8039, -0.9294, -0.9765,  ..., -0.9686, -0.9373, -0.8588],\n",
       "            [-0.9216, -0.9765, -0.9843,  ..., -0.9529, -0.8667, -0.7961]]],\n",
       "  \n",
       "  \n",
       "          [[[ 0.4667,  0.4824,  0.4824,  ...,  0.5137,  0.4902,  0.4745],\n",
       "            [ 0.1922,  0.2078,  0.2392,  ...,  0.5294,  0.5137,  0.4980],\n",
       "            [-0.0980, -0.0824, -0.0431,  ...,  0.5529,  0.5373,  0.5216],\n",
       "            ...,\n",
       "            [-0.2627, -0.2314, -0.3333,  ..., -0.4745, -0.4353, -0.5765],\n",
       "            [-0.2784, -0.2784, -0.3804,  ..., -0.5843, -0.6706, -0.5608],\n",
       "            [-0.3804, -0.4039, -0.4510,  ..., -0.6627, -0.6706, -0.3255]],\n",
       "  \n",
       "           [[ 0.4824,  0.4980,  0.5137,  ...,  0.5451,  0.5216,  0.4980],\n",
       "            [ 0.2235,  0.2392,  0.2863,  ...,  0.5451,  0.5294,  0.5137],\n",
       "            [-0.0118,  0.0118,  0.0588,  ...,  0.5686,  0.5529,  0.5373],\n",
       "            ...,\n",
       "            [ 0.0118,  0.0118, -0.0431,  ..., -0.2078, -0.2706, -0.4745],\n",
       "            [-0.0118, -0.0275, -0.0902,  ..., -0.4118, -0.6000, -0.4196],\n",
       "            [-0.1137, -0.1529, -0.1451,  ..., -0.5216, -0.5765, -0.0824]],\n",
       "  \n",
       "           [[ 0.4353,  0.4510,  0.4588,  ...,  0.5294,  0.5137,  0.4824],\n",
       "            [ 0.2235,  0.2392,  0.2706,  ...,  0.5216,  0.5059,  0.4902],\n",
       "            [-0.0588, -0.0353,  0.0196,  ...,  0.5373,  0.5294,  0.5137],\n",
       "            ...,\n",
       "            [-0.6314, -0.5451, -0.6314,  ..., -0.6392, -0.5843, -0.7255],\n",
       "            [-0.6157, -0.5294, -0.6471,  ..., -0.7333, -0.7961, -0.7647],\n",
       "            [-0.6784, -0.6471, -0.7255,  ..., -0.7647, -0.8118, -0.6471]]],\n",
       "  \n",
       "  \n",
       "          [[[-0.9686, -0.9608, -0.9608,  ...,  0.0667, -0.0275, -0.0431],\n",
       "            [-0.9686, -0.9686, -0.9686,  ...,  0.0510,  0.0118, -0.0118],\n",
       "            [-0.9686, -0.9686, -0.9765,  ...,  0.0824,  0.1294,  0.1216],\n",
       "            ...,\n",
       "            [-0.9529, -0.9451, -0.9529,  ...,  0.8510,  0.7804,  0.7255],\n",
       "            [-0.9294, -0.9373, -0.9451,  ...,  0.8667,  0.7882,  0.7647],\n",
       "            [-0.8510, -0.9294, -0.9529,  ...,  0.9686,  0.9529,  0.9294]],\n",
       "  \n",
       "           [[-0.9686, -0.9608, -0.9608,  ..., -0.9137, -0.9608, -0.9686],\n",
       "            [-0.9686, -0.9686, -0.9686,  ..., -0.9294, -0.9294, -0.9451],\n",
       "            [-0.9686, -0.9686, -0.9765,  ..., -0.9294, -0.8824, -0.8902],\n",
       "            ...,\n",
       "            [-0.9529, -0.9451, -0.9529,  ...,  0.0039, -0.0824, -0.1451],\n",
       "            [-0.9294, -0.9373, -0.9451,  ...,  0.0588, -0.0431, -0.0824],\n",
       "            [-0.8431, -0.9216, -0.9529,  ...,  0.1294,  0.0824,  0.0667]],\n",
       "  \n",
       "           [[-0.9843, -0.9765, -0.9765,  ..., -0.9216, -0.9451, -0.9373],\n",
       "            [-0.9843, -0.9843, -0.9843,  ..., -0.9216, -0.9059, -0.9216],\n",
       "            [-0.9843, -0.9843, -0.9922,  ..., -0.8980, -0.8510, -0.8667],\n",
       "            ...,\n",
       "            [-0.9686, -0.9608, -0.9686,  ..., -0.6549, -0.7098, -0.7490],\n",
       "            [-0.9529, -0.9608, -0.9608,  ..., -0.6157, -0.7333, -0.7725],\n",
       "            [-0.8902, -0.9529, -0.9686,  ..., -0.4980, -0.5451, -0.5451]]],\n",
       "  \n",
       "  \n",
       "          ...,\n",
       "  \n",
       "  \n",
       "          [[[-0.4431, -0.5373, -0.6706,  ..., -0.8196, -0.8510, -0.8196],\n",
       "            [-0.4196, -0.5059, -0.6471,  ..., -0.7412, -0.7490, -0.7176],\n",
       "            [-0.3725, -0.4667, -0.6314,  ..., -0.6392, -0.6392, -0.6471],\n",
       "            ...,\n",
       "            [-0.8196, -0.8667, -0.9059,  ..., -0.6157, -0.6627, -0.6784],\n",
       "            [-0.8039, -0.8745, -0.9137,  ..., -0.6078, -0.6627, -0.6706],\n",
       "            [-0.7569, -0.8431, -0.9059,  ..., -0.6157, -0.6471, -0.6314]],\n",
       "  \n",
       "           [[-0.1059, -0.2471, -0.4510,  ..., -0.8039, -0.8118, -0.7412],\n",
       "            [-0.0745, -0.2078, -0.4196,  ..., -0.6471, -0.6392, -0.5765],\n",
       "            [-0.0196, -0.1529, -0.3804,  ..., -0.4824, -0.4667, -0.4588],\n",
       "            ...,\n",
       "            [-0.6549, -0.7882, -0.8902,  ..., -0.4902, -0.5294, -0.5451],\n",
       "            [-0.6392, -0.7882, -0.8980,  ..., -0.4588, -0.4980, -0.5059],\n",
       "            [-0.5765, -0.7412, -0.8902,  ..., -0.4353, -0.4667, -0.4431]],\n",
       "  \n",
       "           [[-0.5451, -0.6549, -0.8353,  ..., -0.9529, -0.9373, -0.9216],\n",
       "            [-0.5216, -0.6471, -0.8353,  ..., -0.8824, -0.8745, -0.8588],\n",
       "            [-0.4745, -0.6314, -0.8275,  ..., -0.7804, -0.7725, -0.7882],\n",
       "            ...,\n",
       "            [-0.9373, -0.9451, -0.9529,  ..., -0.8588, -0.9059, -0.9059],\n",
       "            [-0.9216, -0.9373, -0.9529,  ..., -0.8275, -0.8824, -0.8745],\n",
       "            [-0.8980, -0.9216, -0.9451,  ..., -0.8118, -0.8510, -0.8196]]],\n",
       "  \n",
       "  \n",
       "          [[[-0.0588, -0.0353, -0.1373,  ..., -0.0510,  0.0745,  0.1529],\n",
       "            [-0.1843, -0.1608, -0.0431,  ...,  0.0980,  0.0039,  0.0118],\n",
       "            [-0.1608, -0.1922, -0.1137,  ...,  0.2078,  0.1373, -0.1608],\n",
       "            ...,\n",
       "            [-0.0196, -0.0902, -0.1529,  ..., -0.4275, -0.4431, -0.3255],\n",
       "            [-0.3961, -0.2863, -0.0745,  ..., -0.3098, -0.2627, -0.3569],\n",
       "            [-0.2863, -0.3569, -0.2314,  ..., -0.3804, -0.2235, -0.3255]],\n",
       "  \n",
       "           [[ 0.2157,  0.2000,  0.0353,  ...,  0.1765,  0.3412,  0.4588],\n",
       "            [ 0.0902,  0.0902,  0.2706,  ...,  0.3804,  0.3020,  0.3098],\n",
       "            [ 0.1608,  0.0667,  0.2000,  ...,  0.4980,  0.4431,  0.1294],\n",
       "            ...,\n",
       "            [ 0.2941,  0.2549,  0.2314,  ..., -0.2157, -0.2392, -0.0902],\n",
       "            [-0.1059, -0.0196,  0.2157,  ..., -0.0039,  0.0196, -0.0196],\n",
       "            [-0.0510, -0.1608, -0.0039,  ..., -0.0902,  0.0431, -0.0275]],\n",
       "  \n",
       "           [[-0.3725, -0.3490, -0.4275,  ..., -0.4118, -0.3020, -0.2706],\n",
       "            [-0.5059, -0.4588, -0.4902,  ..., -0.3412, -0.3961, -0.3961],\n",
       "            [-0.6471, -0.6863, -0.6627,  ..., -0.2549, -0.2314, -0.5686],\n",
       "            ...,\n",
       "            [-0.3882, -0.4745, -0.5451,  ..., -0.7490, -0.8824, -0.8275],\n",
       "            [-0.7255, -0.4745, -0.3098,  ..., -0.6314, -0.6549, -0.7569],\n",
       "            [-0.7490, -0.7255, -0.5765,  ..., -0.7412, -0.5922, -0.6784]]],\n",
       "  \n",
       "  \n",
       "          [[[ 0.0745,  0.0510,  0.0431,  ..., -0.1922,  0.4196,  0.5843],\n",
       "            [ 0.0824,  0.0667,  0.0824,  ...,  0.3961,  0.5765,  0.5451],\n",
       "            [ 0.0667,  0.0431,  0.0745,  ...,  0.5451,  0.3961,  0.2000],\n",
       "            ...,\n",
       "            [-0.6471, -0.6549, -0.6157,  ..., -0.2784, -0.3333, -0.3725],\n",
       "            [-0.6157, -0.5922, -0.5294,  ..., -0.3020, -0.3176, -0.3647],\n",
       "            [-0.5137, -0.5686, -0.5373,  ..., -0.3333, -0.3333, -0.3412]],\n",
       "  \n",
       "           [[ 0.3020,  0.2784,  0.2706,  ..., -0.3412,  0.4745,  0.6627],\n",
       "            [ 0.3098,  0.2941,  0.2863,  ...,  0.4353,  0.6471,  0.6235],\n",
       "            [ 0.2863,  0.2706,  0.2784,  ...,  0.6235,  0.4588,  0.2706],\n",
       "            ...,\n",
       "            [-0.3569, -0.3882, -0.3804,  ..., -0.1686, -0.2000, -0.2235],\n",
       "            [-0.2627, -0.2784, -0.2314,  ..., -0.1922, -0.2000, -0.2314],\n",
       "            [-0.2235, -0.2392, -0.1765,  ..., -0.2549, -0.2157, -0.2235]],\n",
       "  \n",
       "           [[-0.3333, -0.3647, -0.3725,  ..., -0.4196,  0.4275,  0.6314],\n",
       "            [-0.3725, -0.3647, -0.3569,  ...,  0.3882,  0.6078,  0.5843],\n",
       "            [-0.3882, -0.4196, -0.3882,  ...,  0.5608,  0.3725,  0.1529],\n",
       "            ...,\n",
       "            [-0.9765, -0.9765, -0.9373,  ..., -0.6784, -0.7412, -0.7882],\n",
       "            [-0.9686, -0.9529, -0.9529,  ..., -0.7255, -0.7333, -0.7725],\n",
       "            [-0.9373, -0.9765, -0.9765,  ..., -0.7647, -0.7569, -0.7333]]]])})"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loader = torch.utils.data.DataLoader(dataset,\n",
    "                                     batch_size=16,\n",
    "                                     shuffle=True,\n",
    "                                     drop_last=False)\n",
    "\n",
    "len(loader), next(iter(loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c9026717",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "11367.3219"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from diffusers import UNet2DModel\n",
    "\n",
    "#定义模型,随机初始化参数\n",
    "model = UNet2DModel(\n",
    "    sample_size=64,\n",
    "    in_channels=3,\n",
    "    out_channels=3,\n",
    "    layers_per_block=2,\n",
    "    block_out_channels=(128, 128, 256, 256, 512, 512),\n",
    "    down_block_types=(\n",
    "        'DownBlock2D',\n",
    "        'DownBlock2D',\n",
    "        'DownBlock2D',\n",
    "        'DownBlock2D',\n",
    "        'AttnDownBlock2D',\n",
    "        'DownBlock2D',\n",
    "    ),\n",
    "    up_block_types=(\n",
    "        'UpBlock2D',\n",
    "        'AttnUpBlock2D',\n",
    "        'UpBlock2D',\n",
    "        'UpBlock2D',\n",
    "        'UpBlock2D',\n",
    "        'UpBlock2D',\n",
    "    ),\n",
    ")\n",
    "\n",
    "sum(i.numel() for i in model.parameters()) / 10000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2969101e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(DDPMScheduler {\n",
       "   \"_class_name\": \"DDPMScheduler\",\n",
       "   \"_diffusers_version\": \"0.12.1\",\n",
       "   \"beta_end\": 0.02,\n",
       "   \"beta_schedule\": \"linear\",\n",
       "   \"beta_start\": 0.0001,\n",
       "   \"clip_sample\": true,\n",
       "   \"num_train_timesteps\": 1000,\n",
       "   \"prediction_type\": \"epsilon\",\n",
       "   \"trained_betas\": null,\n",
       "   \"variance_type\": \"fixed_small\"\n",
       " },\n",
       " AdamW (\n",
       " Parameter Group 0\n",
       "     amsgrad: False\n",
       "     betas: (0.95, 0.999)\n",
       "     capturable: False\n",
       "     eps: 1e-08\n",
       "     foreach: None\n",
       "     initial_lr: 0.0001\n",
       "     lr: 0.0\n",
       "     maximize: False\n",
       "     weight_decay: 1e-06\n",
       " ),\n",
       " <torch.optim.lr_scheduler.LambdaLR at 0x7f62df1b2af0>,\n",
       " MSELoss())"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from diffusers import DDPMScheduler\n",
    "from diffusers.optimization import get_scheduler\n",
    "\n",
    "#初始化工具类\n",
    "scheduler = DDPMScheduler(num_train_timesteps=1000,\n",
    "                          beta_schedule='linear',\n",
    "                          prediction_type='epsilon')\n",
    "\n",
    "optimizer = torch.optim.AdamW(model.parameters(),\n",
    "                              lr=1e-4,\n",
    "                              betas=(0.95, 0.999),\n",
    "                              weight_decay=1e-6,\n",
    "                              eps=1e-8)\n",
    "\n",
    "scheduler_lr = get_scheduler('cosine',\n",
    "                             optimizer=optimizer,\n",
    "                             num_warmup_steps=500,\n",
    "                             num_training_steps=len(loader) * 100)\n",
    "\n",
    "criterion = torch.nn.MSELoss()\n",
    "\n",
    "scheduler, optimizer, scheduler_lr, criterion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "22436df3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(1.0573, grad_fn=<MseLossBackward0>)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def get_loss(image):\n",
    "    device = image.device\n",
    "    \n",
    "    #随机噪声\n",
    "    #[b, 3, 64, 64]\n",
    "    noise = torch.randn(image.shape).to(device)\n",
    "\n",
    "    #随机b个噪声步数\n",
    "    #1000 = scheduler.config.num_train_timesteps\n",
    "    #[b]\n",
    "    noise_step = torch.randint(0, 1000, (image.shape[0], ),\n",
    "                               device=device).long()\n",
    "\n",
    "    #往图片当中添加噪声\n",
    "    #[b, 3, 64, 64]\n",
    "    image_noise = scheduler.add_noise(image, noise, noise_step)\n",
    "\n",
    "    #把图片里的噪声计算出来\n",
    "    #[b, 3, 64, 64]\n",
    "    out = model(image_noise, noise_step).sample\n",
    "\n",
    "    #求mse loss\n",
    "    return criterion(out, noise)\n",
    "\n",
    "\n",
    "get_loss(torch.randn(16, 3, 64, 64))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "31bd73a7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Cloning https://huggingface.co/lansinuote/diffusion.1.unconditional into local empty directory.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "441eeadfff984242b8d468e13e6ce754",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Download file unet/diffusion_pytorch_model.bin:   0%|          | 15.4k/434M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "717c69fbd52f49c2a2e8d04773bcdb47",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Clean file unet/diffusion_pytorch_model.bin:   0%|          | 1.00k/434M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 87.37168118171394\n",
      "1 22.29192595742643\n",
      "2 19.657229276839644\n",
      "3 17.91914987610653\n",
      "4 17.98647755710408\n",
      "5 17.529722683597356\n",
      "6 18.230531610082835\n",
      "7 16.845244634896517\n",
      "8 17.05660533811897\n",
      "9 16.014264623634517\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5694a0748ced4cba8aeff502d175590f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Upload file unet/diffusion_pytorch_model.bin:   0%|          | 32.0k/434M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "remote: Scanning LFS files for validity...\u001b[K\n",
      "remote: LFS file scan complete.\u001b[K\n",
      "To https://user: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@huggingface.co/lansinuote/diffusion.1.unconditional\n",
      "   c99b906..0f5a8d5  main -> main\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from huggingface_hub import Repository, create_repo\n",
    "from diffusers import DDPMPipeline\n",
    "\n",
    "\n",
    "def train():\n",
    "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "    model.to(device)\n",
    "    model.train()\n",
    "\n",
    "    loss_sum = 0\n",
    "    for epoch in range(10):\n",
    "        for i, data in enumerate(loader):\n",
    "            loss = get_loss(data['image'].to(device))\n",
    "\n",
    "            loss.backward()\n",
    "            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
    "            optimizer.step()\n",
    "            scheduler_lr.step()\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            loss_sum += loss.item()\n",
    "\n",
    "        if epoch % 1 == 0:\n",
    "            print(epoch, loss_sum)\n",
    "            loss_sum = 0\n",
    "\n",
    "    #save\n",
    "    DDPMPipeline(unet=model, scheduler=scheduler).save_pretrained('./save')\n",
    "\n",
    "\n",
    "if push_to_hub:\n",
    "    create_repo(repo_id, exist_ok=True, token=hub_token)\n",
    "    repo = Repository('./save', clone_from=repo_id, token=hub_token)\n",
    "    train()\n",
    "    repo.push_to_hub()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
